import numpy as np


def Obtain_force_torque(w, R, Cm, Ct, Vb, Cd, wb, Cdm, Jrp):
    """
    Function description:
      Obtain the force and torque acting on the body. The airframe of the
      quadrotor is as follows:
    3↓   1↑
      \ /
      / \
    2↑   4↓
    ↑ stands for counterclockwise and ↓stands for clockwise.

    Input：
      w: motor speed，unit：rad/s
      R: body radius(m)
      Cm: torque coefficient
      Ct: thrust coefficient
      Vb: aircraft speed in the body coordinate system
      Cd: damping coefficient
      wb: angular velocity in the body coordinate system trans([p,q,r])
      Cdm: damping moment coefficient vector
      Jrp: moment of inertia of motor rotor + propeller

    Output:
      Fp: thrust under the body coordinate system
      Fd: aerodynamic force under the body coordinate system
      Mp: torque under the body coordinate system
      Md: aerodynamic moment
      Ga: gyro moment
    """
    M_rctcm = np.array([
        [-np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
         -np.sin(np.pi / 4) * R * Ct],  # Roll torque
        [np.sin(np.pi / 4) * R * Ct, -np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
         -np.sin(np.pi / 4) * R * Ct],  # Pitch torque
        [Cm, Cm, -Cm, -Cm]  # Yaw torque
    ])

    Mp = np.dot(M_rctcm, w ** 2)  # Torque

    Fp = np.array([0, 0, -np.sum(Ct * (w ** 2))])  # Thrust
    Fd = -Cd * Vb * np.abs(Vb) * 0.5  # Aerodynamic force
    Md = -Cdm * wb * np.abs(wb)  # Aerodynamic moment

    # Gyro moment
    Ga = np.zeros(3)
    Ga[0] = Jrp * wb[1] * (w[0] + w[1] - w[2] - w[3])
    Ga[1] = Jrp * wb[0] * (-w[0] - w[1] + w[2] + w[3])

    return Fp, Fd, Mp, Md, Ga

def motor_dynamics(w_current, w_target, T, dt):
    """
    电机动力学模型，考虑到电机响应时间的滞后
    w_current: 当前电机速度
    w_target: 目标电机速度（由控制器输出确定）
    T: 电机时间常数
    dt: 时间步长
    """
    return w_current + (w_target - w_current) * dt / T

def uav_dynamics(position, velocity, dcm, rate, w_current,w_target, dt, params):
    """
    简单的无人机动力学模型
    """
    R = params['ModelParam_uavR']
    Cm = params['ModelParam_rotorCm']
    Ct = params['ModelParam_rotorCt']
    Cd = params['ModelParam_uavCd']
    Cdm = params['ModelParam_uavCCm']
    Jrp = params['ModelParam_motorJm']
    g=params['ModelParam_envGravityAcc']
    T = params['ModelParam_motorT']
    #增加电机动力学模型
    w = motor_dynamics(w_current, w_target, T, dt)
    print('motor_rads:',w)
    Fp, Fd, Mp, Md, Ga = Obtain_force_torque(w, R, Cm, Ct, velocity, Cd, rate, Cdm, Jrp)

    mass = params['ModelParam_uavMass']
    J = params['ModelParam_uavJ']
    # print('fdcm:',dcm.dot(np.array([0, 0, mass *g])))
    print('fdcm:', np.array([0, 0, mass * g]).dot(dcm))
    # # 合力计算
    # total_force = Fp + Fd + dcm.dot(np.array([0, 0, mass *g]))
    total_force = Fp + Fd + np.array([0, 0, mass * g]).dot(dcm)
    # acceleration = total_force / mass
    # velocity += acceleration * dt
    # position += velocity * dt

    # 合力矩计算
    total_torque = Mp + Md + Ga
    # angular_acceleration = np.linalg.inv(J).dot(total_torque)
    # rate += angular_acceleration * dt
    # angle += rate * dt

    # return position, velocity, angle, rate
    return total_force,total_torque,w
